import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
import random

# 设置随机种子以保证结果可重复
random.seed(42)
np.random.seed(42)

# 参数设置
GRID_SIZE_X = 50  # 建筑在X轴的网格数
GRID_SIZE_Y = 50  # 建筑在Y轴的网格数
GRID_SIZE_Z = 40  # 建筑在Z轴的网格数（楼层数）

# 楼梯参数
STAIR_WIDTH = 2  # 楼梯宽度
STAIR_LENGTH = 10  # 楼梯长度
STAIR_HEIGHT = 1  # 楼梯高度（每层楼梯的高度）

# 人员参数
NUM_PEOPLE = 100  # 总人数
INITIAL_SPEED = 1.0  # 初始速度（m/s）
A = 0.71  # 速度公式参数
B = 0.51  # 速度公式参数
C_UP = 2.7  # 上楼消耗参数
C_DOWN = 1.0  # 下楼消耗参数

# 时间参数
TIME_STEP = 0.1  # 时间步长（秒）
SIMULATION_TIME = 1000  # 模拟总时间（秒）

# 速度公式
def calculate_speed(v0, distance_sum, time_sum, c):
    """
    计算人员速度
    :param v0: 初始速度
    :param distance_sum: 移动距离累加
    :param time_sum: 时间累加
    :param c: 消耗参数
    :return: 当前速度
    """
    exponent = (distance_sum ** A) / (time_sum ** B)
    speed = v0 * (1 - 0.01 * c * np.exp(exponent) + 0.01)
    return speed

# 元胞自动机类
class Cell:
    def __init__(self, x, y, z, v0=INITIAL_SPEED):
        self.x = x
        self.y = y
        self.z = z
        self.v0 = v0
        self.v = v0
        self.distance_sum = 0
        self.time_sum = 0
        self.time_since_last_move = 0

    def update_speed(self, distance, time_step):
        self.distance_sum += distance
        self.time_sum += time_step
        self.v = calculate_speed(self.v0, self.distance_sum, self.time_sum, C_DOWN if self.z > 0 else C_UP)

    def move(self, dx, dy, dz):
        self.x += dx
        self.y += dy
        self.z += dz
        self.update_speed(np.sqrt(dx**2 + dy**2 + dz**2), TIME_STEP)

# 建筑结构类
class Building:
    def __init__(self, size_x, size_y, size_z):
        self.size_x = size_x
        self.size_y = size_y
        self.size_z = size_z
        # 初始化网格，0为空，1为障碍物，2为人员
        self.grid = np.zeros((size_x, size_y, size_z))
        self.people = []
        self.populate_people()
        self.build_structure()

    def build_structure(self):
        # 添加楼梯，每隔10层一个楼梯
        for z in range(0, self.size_z, 10):
            for x in range(10, 12):
                for y in range(10, 10 + STAIR_LENGTH):
                    self.grid[x, y, z] = 1
            # 添加楼梯连接
            if z + 10 < self.size_z:
                for x in range(10, 12):
                    for y in range(10 + STAIR_LENGTH, 10 + STAIR_LENGTH + 2):
                        self.grid[x, y, z] = 1

    def populate_people(self):
        # 随机分配人员到不同楼层
        for _ in range(NUM_PEOPLE):
            x = random.randint(0, self.size_x - 1)
            y = random.randint(0, self.size_y - 1)
            z = random.randint(0, self.size_z - 1)
            while self.grid[x, y, z] != 0:
                x = random.randint(0, self.size_x - 1)
                y = random.randint(0, self.size_y - 1)
                z = random.randint(0, self.size_z - 1)
            self.grid[x, y, z] = 2
            person = Cell(x, y, z)
            self.people.append(person)

    def get_neighbors(self, x, y, z):
        neighbors = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    nx, ny, nz = x + dx, y + dy, z + dz
                    if 0 <= nx < self.size_x and 0 <= ny < self.size_y and 0 <= nz < self.size_z:
                        if self.grid[nx, ny, nz] == 0:
                            neighbors.append((nx, ny, nz))
        return neighbors

    def update(self, time_step):
        for person in self.people:
            # 获取邻居
            neighbors = self.get_neighbors(person.x, person.y, person.z)
            # 选择最佳移动方向（简单策略：朝向出口）
            # 这里假设出口在z=0层，x=0, y=0
            # 计算距离
            dx = person.x
            dy = person.y
            dz = person.z
            distance = np.sqrt(dx**2 + dy**2 + dz**2)
            # 选择移动方向
            # 简单策略：随机选择一个邻居
            if len(neighbors) > 0:
                nx, ny, nz = random.choice(neighbors)
                # 更新人员位置
                self.grid[person.x, person.y, person.z] = 0
                person.move(nx - person.x, ny - person.y, nz - person.z)
                self.grid[person.x, person.y, person.z] = 2

    def get_positions(self):
        positions = []
        for person in self.people:
            positions.append((person.x, person.y, person.z))
        return positions

# 可视化类
class Visualization:
    def __init__(self, building):
        self.building = building
        self.fig = plt.figure()
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.set_xlim(0, building.size_x)
        self.ax.set_ylim(0, building.size_y)
        self.ax.set_zlim(0, building.size_z)
        self.ax.set_xlabel('X')
        self.ax.set_ylabel('Y')
        self.ax.set_zlabel('Z')
        self.scatter = self.ax.scatter([], [], [], c='blue')

    def update_visualization(self, frame):
        positions = self.building.get_positions()
        x, y, z = zip(*positions)
        self.scatter = self.ax.scatter(x, y, z, c='blue')
        self.ax.set_title(f"Time: {frame * TIME_STEP:.1f} s")
        return self.scatter,

    def animate(self):
        ani = animation.FuncAnimation(self.fig, self.update_visualization, frames=int(SIMULATION_TIME / TIME_STEP), interval=50, blit=True)
        plt.show()

# 主函数
def main():
    building = Building(GRID_SIZE_X, GRID_SIZE_Y, GRID_SIZE_Z)
    visualization = Visualization(building)
    visualization.animate()

if __name__ == "__main__":
    main()